import tushare as ts
import numpy as np
import torch
from model import Model

n_sample = 4096
steps = 60

df = ts.get_hist_data("002415", start="2020-01-01")[::-1]
x = np.float32(df[["open", "high", "low", "close", "volume"]])
x = np.log(x[np.all(x > 0, 1)])
base = x[:-1].copy()
base[..., :4] = base[..., 3:4]
x = x[1:] - base

MEAN, STD = map(np.load("tr.npz").get, ["mean", "std"])

x_data = (x[None, ...] - MEAN) / STD
x_data = torch.from_numpy(x_data)

m = Model()
m.load_state_dict(torch.load("model.pt", map_location='cpu'))
dist, h = m(x_data[:, :-1], ret_h=True)
h = (h[0].repeat(1, n_sample, 1), h[1].repeat(1, n_sample, 1))
dist, h = m(x_data[:, -1:].repeat(n_sample, 1, 1), h, ret_h=True)

samples = []
_prev = np.float32(df["close"].iloc[-1])
totvalue = torch.tensor(_prev).repeat(n_sample)
for step in range(steps):
    x_pred = dist.sample()[..., None]
    totvalue *= (x_pred.flatten() * STD[3] + MEAN[3]).exp()
    samples.append(totvalue.clone())
    dist, h = m(x_pred, h, ret_h=True)
samples = torch.stack(samples)
mean = samples.mean(1)
ci0 = np.percentile(samples, 95, 1)
ci1 = np.percentile(samples, 5, 1)
ci2 = np.percentile(samples, 75, 1)
ci3 = np.percentile(samples, 25, 1)
# plt.yscale('log')
plt.plot(range(-1, steps), [_prev, *mean], label="$\mathbb{E}(\cdot)$")
if n_sample > 1:
    plt.plot(range(-1, steps), [_prev, *ci0], c='gray', ls=('dotted'), label=".9CI")
    plt.plot(range(-1, steps), [_prev, *ci1], c='gray', ls=('dotted'))
    plt.plot(range(-1, steps), [_prev, *ci2], c='lightblue', ls=('dotted'), label=".5CI")
    plt.plot(range(-1, steps), [_prev, *ci3], c='lightblue', ls=('dotted'))
plt.legend()

# prev
plt.plot(range(-steps // 2, 0), raw_price["totvalue"][-steps // 2:])

# gt
plt.title(f"Expected annualized return: {((mean[-1].item() / _prev) ** (200 / steps) - 1) * 100: .2f}%")
plt.xlabel("Days")
plt.ylabel("SH000300")
plt.savefig("plt.jpg")
plt.close()

